package com.ruoyi.web.controller.tool;

import com.google.gson.JsonObject;
import com.ruoyi.common.core.domain.AjaxResult;
import com.ruoyi.common.utils.MilvusUtils;
import io.milvus.v2.service.vector.response.SearchResp;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@RestController
@RequestMapping("/milvusTest")
public class MilvusTest {


    @GetMapping("/createPartition")
    public AjaxResult createPartition() {
        MilvusUtils.createPartition("test", "Mexico");
        return AjaxResult.success();
    }

    @GetMapping("/insertData")
    public AjaxResult insertData() {

        List<JsonObject> jsonObjects = new ArrayList<>();
        JsonObject jsonObject = new JsonObject();
        jsonObject.addProperty("text", "小米空调，这是一款来自小米的空调制冷效果好");
        jsonObjects.add(jsonObject);

        JsonObject jsonObject1 = new JsonObject();
        jsonObject1.addProperty("text", "小米冰箱，小米冰箱省电");
        jsonObjects.add(jsonObject1);

        JsonObject jsonObject2 = new JsonObject();
        jsonObject2.addProperty("text", "海尔冰箱，价格优惠");
        jsonObjects.add(jsonObject2);

        JsonObject jsonObject3 = new JsonObject();
        jsonObject3.addProperty("text", "美的电风扇，夏天必备神器");
        jsonObjects.add(jsonObject3);

        boolean test = MilvusUtils.insertData("test", "test", jsonObjects);
        return AjaxResult.success(test);
    }

    @GetMapping("/search")
    public AjaxResult search(@RequestParam("query") String query) {
        SearchResp searchResp = MilvusUtils.sparseSearchResp("test", Collections.singletonList("test"), query, 10, Collections.singletonList("text"));
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        return AjaxResult.success(searchResults);
    }

}
